#!/usr/bin/env python

# https://gitee.com/williamzjc/mcp-llm/

import asyncio

from mcp.client.stdio import stdio_client
from mcp import ClientSession, StdioServerParameters

# Set the list of tools
tools = [
        {
            "type": "function",
            "description": "Addition operation",
            "function": {
                "name": "add",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "a": {"type": "number"},
                        "b": {"type": "number"}
                    },
                    "required": ["a", "b"]
                }
            }
        },
        {
            "type": "function",
            "description": "Multiplication operation",
            "function": {
                "name": "multiply",
                "parameters": {
                    "type": "object",
                    "properties": {
                        "a": {"type": "number"},
                        "b": {"type": "number"}
                    },
                    "required": ["a", "b"]
                }
            }
        }
    ]


tool_names = [tool["function"]["name"] for tool in tools]


params = StdioServerParameters(
    command="python3",
    args=["server.py"])


async def main():

    async with stdio_client(params) as (stdio, write):
        async with ClientSession(stdio, write) as session:

            await session.initialize()

            # Get the name of the function and its arguments
            fn_name = 'add'
            fn_args = {"a":1, "b":3}

            if fn_name in tool_names:
                # Call the function and get the result
                result = await session.call_tool(fn_name, fn_args)
                print(f"System: Model decides to call tool: {fn_name}, the result is {result.content[0].text}.")
            else:
                print("Function not found!")


if __name__ == "__main__":
    asyncio.run(main())
